
import JT
import time

# Define which data to use for registration:
vert = 0
proj = 1
body = 1

Data = JT.Data.gold_standard() 
#===============================================================================
#    Setup DRR
#===============================================================================      
drr = JT.DRR()
drr.SetExternalCalibration(Data.Frames[vert].FixedImages[proj].calibration)

drr.AddVolume(Data.Frames[vert].Volumes[0].data)
drr._renWin.SetOffScreenRendering(0)

reg=JT.Registration()
reg.SetFixedImage(Data.Frames[vert].FixedImages[proj].data.GetImage(0))
reg.SetMovingImageGenerator(drr)
reg.SetRegionOfInterest(Data.Frames[vert].FixedImages[proj].ROIs[body].data)

Testing = JT.Testing()

#===============================================================================
#    Setup Registration
#===============================================================================      
# set Metric
reg.SetImageMetric('bpc')
# Set Optimizer
gd = JT.RegularStepGradientDescentOptimizer()
reg.SetOptimizer(gd)

# Set Scales 1/1000 for translations, 1 for rotations
scales = [1.0,1.0,1.0,1700.0,1700.0,1700.0]
reg.SetScales(scales)

# Registration Options
reg._plot = False
reg._debug = True
reg._logGradientResult = False
gd._debug = True

#===============================================================================
# Setup Optimizer Parameters
#===============================================================================
# Set the cost function
gd.SetCostFunction(reg)

# Modify Optimizer Parameters
gd.SetMaximumStepLength(10)
gd.SetMinimumStepLength(.5)
gd.SetRelaxationFactor(0.75)
gd.SetNumberOfIterations(50)
gd._MaximumGradientStepLength = 4.0
gd._MinimumGradientStepLength = 0.25

def run(n_pose=100):    
    # Set the starting pose
    pose = Testing.GetStartPose(n_pose)
    print "Starting Pose: ", pose
    reg.SetInitialPose(pose)
    # Start the registration process
    reg.Start()
    # Print Final State    
    gd.PrintSelf()
    Data.Frames[vert].FixedImages[proj].ROIs[body].Pose.SetPose(reg._BestPose)
        

if __name__ == "__main__":
    #T = Test(logging=False,upload=False)
    run()
    time.sleep(7)
    JT.PlotResultsFromFile.main(0)
    
